import os, sys, cv2
import numpy as np
from PySide6.QtWidgets import QApplication, QWidget
from PySide6.QtGui import QPainter, QColor, QPen, QGuiApplication, QPixmap, QImage
from PySide6.QtCore import Qt, QRect


def qpixmap_to_cv2(pixmap):
    qimg = pixmap.toImage()
    buffer_ = qimg.bits()
    buffer_.cast('B').release()
    cv2_img = np.asarray(buffer_).reshape(qimg.height(), qimg.width(), 4)
    return cv2_img

class ScreenshotOverlay(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        # 半透明黑色背景
        self.setAttribute(Qt.WA_TranslucentBackground, True)
        self.screen = QGuiApplication.primaryScreen()
        self.full_pixmap = self.screen.grabWindow(0)
        self.origin = None
        self.current = None
        self.selection = QRect()

    def showEvent(self, event):
        self.resize(self.screen.size())
        self.showFullScreen()

    def paintEvent(self, event):
        painter = QPainter(self)
        # 绘制截图背景
        painter.drawPixmap(0, 0, self.full_pixmap)
        # 蒙层
        painter.fillRect(self.rect(), QColor(0, 0, 0, 100))
        # 清除选区蒙层
        if not self.selection.isNull():
            painter.setCompositionMode(QPainter.CompositionMode_Clear)
            painter.fillRect(self.selection, QColor(0,0,0,0))
            # 画选区边框
            painter.setCompositionMode(QPainter.CompositionMode_SourceOver)
            pen = QPen(QColor(255,255,255), 2)
            painter.setPen(pen)
            painter.drawRect(self.selection)
        painter.end()

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.origin = event.pos()
            self.selection = QRect(self.origin, self.origin)

    def mouseMoveEvent(self, event):
        if self.origin:
            self.current = event.pos()
            self.selection = QRect(self.origin, self.current).normalized()
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and not self.selection.isNull():
            # 截取选区
            cropped = self.full_pixmap.copy(self.selection)
            cvimg = qpixmap_to_cv2(cropped)
            cv2.imshow("cropped", cvimg)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
            # 进入注释画布
            # self.open_annotation(cropped)
            self.close()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setStyle('Fusion')
    win = ScreenshotOverlay()
    win.show()
    sys.exit(app.exec())
